# main.py
import argparse
from utils.train import train
from utils.inference import inference
from config import Config
import torch
import os

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Neural Machine Translation with Transformer')
    parser.add_argument('--mode', type=str, choices=['train', 'inference'], default='train',
                        help='运行模式: train or inference (default: train)')
    
    args = parser.parse_args()
    
    # 加载配置
    config = Config()
       
    # 打印配置信息
    print("运行配置:")
    print(f"  模式: {args.mode}")
    print(f"  设备: {config.device}")
    print(f"  批次大小: {config.batch_size}")
    if args.mode == 'train':
        print(f"  训练轮数: {config.epochs}")
        print(f"  学习率: {config.learning_rate}")
    
    # 根据模式运行
    if args.mode == 'train':
        train(config)
    else:
        inference(config)

if __name__ == '__main__':
    main()